import math
import heapq
# 存储点的集合和目标点
points = []
targetPoint = ()

# 定义节点类
class Node:
    # 初始化节点，节点有四个属性：
    # point代表当前节点的点,
    # leftPlane代表划分的左平面,
    # rightPlane代表划分的右平面,
    # dimension代表当前划分的维度
    def __init__(self, point, leftPlane, rightPlane, dimension):
        self.point = point
        self.leftPlane = leftPlane
        self.rightPlane = rightPlane
        self.dimension = dimension

# 定义KDTree类
class KDTree:
    # 初始化KDTree
    def __init__(self, data):
        self.shortestDistancePoint = None  # 存储最近点
        self.shortestDistanceInPoints = math.inf  # 存储最近距离
        self.points = data  # 点的集合

    # 生成KDTree
    def generateKDTree(self, pts, dimension):
        if len(pts) == 0:
            return None  # 当前平面的点为空就结束
        midPoint = self.getMidPoint(pts)  # 获取当前最中间的点
        sortedPts = self.sortByDimension(pts, dimension)  # 按当前维度排序
        # 递归生成左平面
        leftPlane = self.generateKDTree(sortedPts[:midPoint], self.nextDimension(dimension))
        # 递归生成右平面
        rightPlane = self.generateKDTree(sortedPts[midPoint + 1:], self.nextDimension(dimension))
        return Node(sortedPts[midPoint], leftPlane, rightPlane, dimension)

    # 获取中间节点
    @staticmethod
    def getMidPoint(pts):
        return len(pts) // 2

    # 按维度排序
    @staticmethod
    def sortByDimension(pts, dimension):
        return sorted(pts, key=lambda point: point[dimension])

    # 下一维度（循环）
    @staticmethod
    def nextDimension(dimension):
        return (dimension + 1) % len(targetPoint)

    # k近邻搜索
    def knnSearch(self, node, targetPt, k, heap):
        if node is None:
            return
        # 计算当前节点与目标点在划分维度上的距离
        dist = node.point[node.dimension] - targetPt[node.dimension]
        # 根据距离选择要递归搜索的平面
        if dist > 0:  # 目标点在节点的左侧或上侧
            self.knnSearch(node.leftPlane, targetPt, k, heap)
        else:  # 目标点在节点的右侧或下侧
            self.knnSearch(node.rightPlane, targetPt, k, heap)

        # 计算当前节点与目标点的距离
        tempDis = self.getDistance(node.point, targetPt)
        # 如果优先队列未满，则将当前节点加入优先队列
        if len(heap) < k:
            heapq.heappush(heap, (-tempDis, node.point))
        # 如果优先队列已满且当前节点距离更近，则替换队列中最远的点
        elif tempDis < -heap[0][0]:
            heapq.heappushpop(heap, (-tempDis, node.point))

        # 回溯：检查另一侧平面是否有可能存在更近的点
        if abs(dist) < -heap[0][0] or len(heap) < k:
            if dist > 0:
                self.knnSearch(node.rightPlane, targetPt, k, heap)
            else:
                self.knnSearch(node.leftPlane, targetPt, k, heap)

    # 计算两个点之间的距离
    @staticmethod
    def getDistance(p0, p1):
        distanceSum = 0
        for i in range(len(targetPoint)):
            distanceSum += (p0[i] - p1[i]) ** 2
        return math.sqrt(distanceSum)

    # 获取最近的k个点
    def getKNearestPts(self, root, targetPt, k):
        heap = []  # 优先队列，存储最近的k个点
        self.knnSearch(root, targetPt, k, heap)
        return sorted([(-dist, pt) for dist, pt in heap])

if __name__ == '__main__':
    points = [(2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2)]
    targetPoint = (9, 2)
    kdtree = KDTree(points)
    treeRoot = kdtree.generateKDTree(points, 0)
    k = 3
    nearestPts = kdtree.getKNearestPts(treeRoot, targetPoint, k)
    print(f"最近的{len(nearestPts)}个点是:")
    for dist, pt in nearestPts:
        print(f"点: {pt}, 距离: {dist}")
